{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/tomasonjo/blogs/blob/master/diffbot/diffbot_import.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RQ5c2LRxUknp",
    "outputId": "3cc2a4ee-e3dd-4ac8-87fd-d443e1b58974"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting neo4j\n",
      "  Downloading neo4j-5.24.0-py3-none-any.whl.metadata (5.7 kB)\n",
      "Requirement already satisfied: pytz in /usr/local/lib/python3.10/dist-packages (from neo4j) (2024.1)\n",
      "Downloading neo4j-5.24.0-py3-none-any.whl (294 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m294.6/294.6 kB\u001b[0m \u001b[31m3.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: neo4j\n",
      "Successfully installed neo4j-5.24.0\n"
     ]
    }
   ],
   "source": [
    "!pip install neo4j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "9cR3kuPRCmXU"
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "token = \"\"\n",
    "data = requests.get(f\"https://kg.diffbot.com/kg/v3/dql?type=query&token={token}&query=id%3AE9nQaOvC9MXaUzXCuN81tFQ&size=1\").json()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "hraRRIc_Cvt6",
    "outputId": "57cf1480-1203-4812-ca0f-0f18c74c5a11"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['twitterUri', 'nbActiveEmployeeEdges', 'suppliers', 'type', 'naceClassification', 'iSicClassification', 'allNames', 'isAcquired', 'revenue', 'competitors', 'emailAddresses', 'yearlyRevenues', 'logo', 'id', 'stock', 'nbOrigins', 'sicClassification', 'foundingDate', 'image', 'images', 'wikipediaUri', 'secForms', 'irsEmployerIdentificationNumbers', 'naicsClassification2017', 'diffbotUri', 'nbIncomingEdges', 'nbEmployeesMin', 'ipo', 'employeeCategories', 'wikipediaPageviewsLastQuarterGrowth', 'capitalization', 'wikipediaPageviewsLastYear', 'angellistUri', 'name', 'legalEntities', 'nbEmployeesMax', 'totalInvestment', 'allOriginHashes', 'linkedInUri', 'naicsClassification', 'nbEmployees', 'monthlyTraffic', 'githubUri', 'isDissolved', 'companiesHouseIds', 'importance', 'origin', 'description', 'homepageUri', 'founders', 'ceo', 'investments', 'phoneNumbers', 'diffbotClassification', 'blogUri', 'descriptors', 'wikipediaPageviewsLastYearGrowth', 'partnerships', 'isNonProfit', 'japaneseCorporateNumbers', 'origins', 'isPublic', 'wikipediaPageviews', 'categories', 'customers', 'crawlTimestamp', 'nbUniqueInvestors', 'facebookUri', 'secCentralIndexKeys', 'summary', 'types', 'technographics', 'boardMembers', 'wikipediaPageviewsLastQuarter', 'fullName', 'allUris', 'remoteScore', 'monthlyTrafficGrowth', 'quarterlyRevenues', 'nbLocations', 'crunchbaseUri', 'googlePlusUri', 'industries', 'allDescriptions', 'location', 'locations', 'subsidiaries', 'wikipediaPageviewsGrowth', 'vatIdentificationNumbers'])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['data'][0]['entity'].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9gcxoW44ELwY",
    "outputId": "15fa9544-9c8e-4653-d59c-a99873fba104"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "twitterUri <class 'str'>\n",
      "nbActiveEmployeeEdges <class 'int'>\n",
      "suppliers <class 'list'>\n",
      "type <class 'str'>\n",
      "naceClassification <class 'list'>\n",
      "iSicClassification <class 'list'>\n",
      "allNames <class 'list'>\n",
      "isAcquired <class 'bool'>\n",
      "revenue <class 'dict'>\n",
      "competitors <class 'list'>\n",
      "emailAddresses <class 'list'>\n",
      "yearlyRevenues <class 'list'>\n",
      "logo <class 'str'>\n",
      "id <class 'str'>\n",
      "stock <class 'dict'>\n",
      "nbOrigins <class 'int'>\n",
      "sicClassification <class 'list'>\n",
      "foundingDate <class 'dict'>\n",
      "image <class 'str'>\n",
      "images <class 'list'>\n",
      "wikipediaUri <class 'str'>\n",
      "secForms <class 'list'>\n",
      "irsEmployerIdentificationNumbers <class 'list'>\n",
      "naicsClassification2017 <class 'list'>\n",
      "diffbotUri <class 'str'>\n",
      "nbIncomingEdges <class 'int'>\n",
      "nbEmployeesMin <class 'int'>\n",
      "ipo <class 'dict'>\n",
      "employeeCategories <class 'list'>\n",
      "wikipediaPageviewsLastQuarterGrowth <class 'float'>\n",
      "capitalization <class 'dict'>\n",
      "wikipediaPageviewsLastYear <class 'int'>\n",
      "angellistUri <class 'str'>\n",
      "name <class 'str'>\n",
      "legalEntities <class 'list'>\n",
      "nbEmployeesMax <class 'int'>\n",
      "totalInvestment <class 'dict'>\n",
      "allOriginHashes <class 'list'>\n",
      "linkedInUri <class 'str'>\n",
      "naicsClassification <class 'list'>\n",
      "nbEmployees <class 'int'>\n",
      "monthlyTraffic <class 'int'>\n",
      "githubUri <class 'str'>\n",
      "isDissolved <class 'bool'>\n",
      "companiesHouseIds <class 'list'>\n",
      "importance <class 'float'>\n",
      "origin <class 'str'>\n",
      "description <class 'str'>\n",
      "homepageUri <class 'str'>\n",
      "founders <class 'list'>\n",
      "ceo <class 'dict'>\n",
      "investments <class 'list'>\n",
      "phoneNumbers <class 'list'>\n",
      "diffbotClassification <class 'list'>\n",
      "blogUri <class 'str'>\n",
      "descriptors <class 'list'>\n",
      "wikipediaPageviewsLastYearGrowth <class 'float'>\n",
      "partnerships <class 'list'>\n",
      "isNonProfit <class 'bool'>\n",
      "japaneseCorporateNumbers <class 'list'>\n",
      "origins <class 'list'>\n",
      "isPublic <class 'bool'>\n",
      "wikipediaPageviews <class 'int'>\n",
      "categories <class 'list'>\n",
      "customers <class 'list'>\n",
      "crawlTimestamp <class 'int'>\n",
      "nbUniqueInvestors <class 'int'>\n",
      "facebookUri <class 'str'>\n",
      "secCentralIndexKeys <class 'list'>\n",
      "summary <class 'str'>\n",
      "types <class 'list'>\n",
      "technographics <class 'list'>\n",
      "boardMembers <class 'list'>\n",
      "wikipediaPageviewsLastQuarter <class 'int'>\n",
      "fullName <class 'str'>\n",
      "allUris <class 'list'>\n",
      "remoteScore <class 'float'>\n",
      "monthlyTrafficGrowth <class 'float'>\n",
      "quarterlyRevenues <class 'list'>\n",
      "nbLocations <class 'int'>\n",
      "crunchbaseUri <class 'str'>\n",
      "googlePlusUri <class 'str'>\n",
      "industries <class 'list'>\n",
      "allDescriptions <class 'list'>\n",
      "location <class 'dict'>\n",
      "locations <class 'list'>\n",
      "subsidiaries <class 'list'>\n",
      "wikipediaPageviewsGrowth <class 'float'>\n",
      "vatIdentificationNumbers <class 'list'>\n"
     ]
    }
   ],
   "source": [
    "for key in data['data'][0]['entity'].keys():\n",
    "    print(key, type(data['data'][0]['entity'][key]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qAWzw6nBFUlY",
    "outputId": "9f060429-fa9d-4e44-d0c7-cf43da55db12"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'revenue': {'currency': 'USD', 'value': 60922000000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2024},\n",
       " {'revenue': {'currency': 'USD', 'value': 26974000000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2023},\n",
       " {'revenue': {'currency': 'USD', 'value': 26914000000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2022},\n",
       " {'revenue': {'currency': 'USD', 'value': 16675000000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2021},\n",
       " {'revenue': {'currency': 'USD', 'value': 10918000000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2020},\n",
       " {'revenue': {'currency': 'USD', 'value': 11716000000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2019},\n",
       " {'revenue': {'currency': 'USD', 'value': 9714000000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2018},\n",
       " {'revenue': {'currency': 'USD', 'value': 9714000000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2017},\n",
       " {'revenue': {'currency': 'USD', 'value': 6910000000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2016},\n",
       " {'revenue': {'currency': 'USD', 'value': 5010000000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2015},\n",
       " {'revenue': {'currency': 'USD', 'value': 4682000000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2014},\n",
       " {'revenue': {'currency': 'USD', 'value': 4130000000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2013},\n",
       " {'revenue': {'currency': 'USD', 'value': 4280159000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2012},\n",
       " {'revenue': {'currency': 'USD', 'value': 3997930000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2011},\n",
       " {'revenue': {'currency': 'USD', 'value': 3543309000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2010},\n",
       " {'revenue': {'currency': 'USD', 'value': 3326445000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2009},\n",
       " {'revenue': {'currency': 'USD', 'value': 3424859000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2008},\n",
       " {'revenue': {'currency': 'USD', 'value': 4097860000.0},\n",
       "  'isCurrent': False,\n",
       "  'year': 2007}]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['data'][0]['entity']['yearlyRevenues']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "to5LtrRGE9yU",
    "outputId": "11b685f6-241b-46ef-e628-569ac2709720"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['yearlyRevenues', 'quarterlyRevenues']\n"
     ]
    }
   ],
   "source": [
    "used = [\"diffbotUri\", \"name\", \"summary\", \"wikipediaUri\", \"fullName\", \"twitterUri\", \"isAcquired\", \"linkedInUri\", \"githubUri\", \"importance\", \"blogUri\", \"homepageUri\",\n",
    "        \"angellistUri\", \"logo\", \"monthlyTraffic\", \"isNonProfit\", \"nbEmployees\", \"crunchbaseUri\", \"googlePlusUri\", \"description\", \"suppliers\", \"foundingDate\", \"competitors\", \"stock\",\n",
    "        \"isPublic\", \"wikipediaPageviews\", \"nbUniqueInvestors\", \"totalInvestment\", \"isDissolved\", \"facebookUri\", \"founders\", \"partnerships\", \"customers\", \"boardMembers\", \"subsidiaries\",\n",
    "        \"diffbotClassification\", \"companiesHouseIds\", \"secCentralIndexKeys\", \"ceo\", \"revenue\", \"ipo\", \"capitalization\", \"locations\", \"technographics\", \"foundingDateClean\", \"ipoDate\",\n",
    "        \"secForms\", \"investments\"]\n",
    "ignore = [\"origin\", \"type\", \"nbActiveEmployeeEdges\", \"id\", \"naceClassification\", \"nbOrigins\", \"iSicClassification\", \"allNames\", \"emailAddresses\", \"nbIncomingEdges\", \"nbEmployeesMin\",\n",
    "          \"wikipediaPageviewsLastQuarterGrowth\", \"japaneseCorporateNumbers\", \"wikipediaPageviewsLastYear\", \"crawlTimestamp\", \"allUris\", \"remoteScore\", \"monthlyTrafficGrowth\", \"nbLocations\",\n",
    "          \"wikipediaPageviewsGrowth\", \"vatIdentificationNumbers\", \"image\", \"images\", \"nbEmployeesMax\", \"sicClassification\", \"naicsClassification2017\", \"allOriginHashes\", \"phoneNumbers\",\n",
    "          \"wikipediaPageviewsLastYearGrowth\", \"origins\", \"wikipediaPageviewsLastQuarter\", \"allDescriptions\", \"naicsClassification\", \"employeeCategories\", \"categories\", \"types\",\n",
    "          \"industries\", \"location\", \"irsEmployerIdentificationNumbers\", \"legalEntities\", \"descriptors\"]\n",
    "available = [key for key in data['data'][0]['entity'].keys() if key not in used and key not in ignore]\n",
    "print(available)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "OCiiUmWrUrig"
   },
   "outputs": [],
   "source": [
    "from neo4j import GraphDatabase\n",
    "url = \"bolt://localhost:7687\"\n",
    "user = \"neo4j\"\n",
    "password = \"password\"\n",
    "\n",
    "driver = GraphDatabase.driver(url, auth=(user, password))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "sEdler9AbJew"
   },
   "outputs": [],
   "source": [
    "from typing import Optional, Union\n",
    "from datetime import datetime\n",
    "\n",
    "def get_datetime(value: Optional[Union[str, int, float]]) -> datetime:\n",
    "    if not value:\n",
    "        return value\n",
    "    return datetime.fromtimestamp(float(value) / 1000.0)\n",
    "\n",
    "def preprocess_dates(data):\n",
    "    for entity in data['data']:\n",
    "        entity['entity']['foundingDateClean'] = get_datetime(entity['entity']['foundingDate']['timestamp'])\n",
    "        entity['entity']['ipoDate'] = get_datetime(entity['entity']['ipo']['date']['timestamp'])\n",
    "        for sec in entity['entity']['secForms']:\n",
    "            sec['dateClean'] = get_datetime(sec['filingDate']['timestamp'])\n",
    "        for inv in entity['entity']['investments']:\n",
    "            inv['dateClean'] = get_datetime(inv['date']['timestamp'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "zfaAFIiNdtd6"
   },
   "outputs": [],
   "source": [
    "preprocess_dates(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "1ZPAuesYDcGJ"
   },
   "outputs": [],
   "source": [
    "driver.execute_query(\"CREATE CONSTRAINT IF NOT EXISTS FOR (n:Organization) REQUIRE n.id IS UNIQUE;\")\n",
    "driver.execute_query(\"CREATE CONSTRAINT IF NOT EXISTS FOR (n:Person) REQUIRE n.id IS UNIQUE;\")\n",
    "driver.execute_query(\"CREATE CONSTRAINT IF NOT EXISTS FOR (n:Classification) REQUIRE n.id IS UNIQUE;\")\n",
    "driver.execute_query(\"CREATE CONSTRAINT IF NOT EXISTS FOR (n:Place) REQUIRE n.address IS UNIQUE;\")\n",
    "driver.execute_query(\"CREATE CONSTRAINT IF NOT EXISTS FOR (n:Country) REQUIRE n.id IS UNIQUE;\")\n",
    "driver.execute_query(\"CREATE CONSTRAINT IF NOT EXISTS FOR (n:Technographic) REQUIRE n.name IS UNIQUE;\")\n",
    "driver.execute_query(\"CREATE CONSTRAINT IF NOT EXISTS FOR (n:TechnographicCategory) REQUIRE n.name IS UNIQUE;\")\n",
    "driver.execute_query(\"CREATE CONSTRAINT IF NOT EXISTS FOR (n:SecForm) REQUIRE n.filingUrl IS UNIQUE;\")\n",
    "driver.execute_query(\"CREATE CONSTRAINT IF NOT EXISTS FOR (n:InvestmentSeries) REQUIRE n.id IS UNIQUE;\")\n",
    "driver.execute_query(\"CREATE CONSTRAINT IF NOT EXISTS FOR (n:RevenueYear) REQUIRE n.id IS UNIQUE;\")\n",
    "\n",
    "\n",
    "query = \"\"\"\n",
    "UNWIND $data AS row\n",
    "MERGE (o:Organization {id: row.diffbotUri})\n",
    "SET o.name = row.name,\n",
    "    o.summary = row.summary,\n",
    "    o.description = row.description,\n",
    "    o.wikipediaUri = row.wikipediaUri,\n",
    "    o.fullName = row.fullName,\n",
    "    o.twitterUri = row.twitterUri,\n",
    "    o.isAcquired = row.isAcquired,\n",
    "    o.isPublic = row.isPublic,\n",
    "    o.isDissolved = row.isDissolved,\n",
    "    o.isNonProfit = row.isNonProfit,\n",
    "    o.linkedInUri = row.linkedInUri,\n",
    "    o.crunchbaseUri = row.crunchbaseUri,\n",
    "    o.googlePlusUri = row.googlePlusUri,\n",
    "    o.githubUri = row.githubUri,\n",
    "    o.blogUri = row.blogUri,\n",
    "    o.facebookUri = row.facebookUri,\n",
    "    o.homepageUri = row.homepageUri,\n",
    "    o.angellistUri = row.angellistUri,\n",
    "    o.companiesHouseIds = row.companiesHouseIds,\n",
    "    o.secCentralIndexKeys = row.secCentralIndexKeys,\n",
    "    o.logo = row.logo,\n",
    "    o.importance = toFloat(row.importance),\n",
    "    o.monthlyTraffic = toFloat(row.monthlyTraffic),\n",
    "    o.nbEmployees = toInteger(row.nbEmployees),\n",
    "    o.ipoStockExchange = row.ipo.stockExchange,\n",
    "    o.ipoDate = row.ipoDate,\n",
    "    o.foundingDate = row.foundingDateClean,\n",
    "    o.stockTicker = CASE WHEN row.stock IS NOT NULL THEN row.stock.ticker ELSE null END,\n",
    "    o.wikipediaPageviews = toInteger(row.wikipediaPageviews),\n",
    "    o.nbUniqueInvestors = toInteger(row.nbUniqueInvestors),\n",
    "    o.totalInvestment = toFloat(row.totalInvestment.value),\n",
    "    o.totalInvestmentCurrency = row.totalInvestment.currency,\n",
    "    o.revenue = toFloat(row.revenue.value),\n",
    "    o.revenueCurrency = row.revenue.currency,\n",
    "    o.capitalization = toFloat(row.capitalization.value),\n",
    "    o.capitalizationCurrency = row.capitalization.currency\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    UNWIND [x in row.suppliers WHERE NOT x.diffbotUri is null] AS supplier\n",
    "    MERGE (s:`Organization` {id: supplier.diffbotUri})\n",
    "    ON CREATE SET s.name = supplier.name,\n",
    "                  s.summary = supplier.summary\n",
    "    MERGE (o)-[:HAS_SUPPLIER]->(s)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    UNWIND [x in row.competitors WHERE NOT x.diffbotUri is null] AS competitor\n",
    "    MERGE (s:`Organization` {id: competitor.diffbotUri})\n",
    "    ON CREATE SET s.name = competitor.name,\n",
    "                  s.summary = competitor.summary\n",
    "    MERGE (o)-[:HAS_COMPETITOR]->(s)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    UNWIND [x in row.founders WHERE NOT x.diffbotUri is null] AS founder\n",
    "    MERGE (s:`Person` {id: founder.diffbotUri})\n",
    "    ON CREATE SET s.name = founder.name,\n",
    "                  s.summary = founder.summary\n",
    "    MERGE (o)-[:HAS_FOUNDER]->(s)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    UNWIND [x in row.partnerships WHERE NOT x.diffbotUri is null] AS partnership\n",
    "    MERGE (s:`Organization` {id: partnership.diffbotUri})\n",
    "    ON CREATE SET s.name = partnership.name,\n",
    "                  s.summary = partnership.summary\n",
    "    MERGE (o)-[:PARTNERSHIP]-(s)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    UNWIND [x in row.customers WHERE NOT x.diffbotUri is null] AS customer\n",
    "    MERGE (s:`Organization` {id: customer.diffbotUri})\n",
    "    ON CREATE SET s.name = customer.name,\n",
    "                  s.summary = customer.summary\n",
    "    MERGE (o)-[:HAS_CUSTOMER]->(s)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    UNWIND [x in row.boardMembers WHERE NOT x.diffbotUri is null] AS board_member\n",
    "    MERGE (s:`Person` {id: board_member.diffbotUri})\n",
    "    ON CREATE SET s.name = board_member.name,\n",
    "                  s.summary = board_member.summary\n",
    "    MERGE (o)-[:BOARD_MEMBER]->(s)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    UNWIND [x in row.subsidiaries WHERE NOT x.diffbotUri is null] AS subsidiary\n",
    "    MERGE (s:`Organization` {id: subsidiary.diffbotUri})\n",
    "    ON CREATE SET s.name = subsidiary.name,\n",
    "                  s.summary = subsidiary.summary\n",
    "    MERGE (o)-[:HAS_SUBSIDIARY]->(s)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    UNWIND row.secForms AS secForm\n",
    "    MERGE (s:`SecForm` {filingUrl: secForm.filingUrl})\n",
    "    ON CREATE SET s.formType = secForm.formType,\n",
    "                  s.documentUrl = secForm.documentUrl,\n",
    "                  s.filingDate = secForm.dateClean\n",
    "    MERGE (o)-[:HAS_SEC_FORM]->(s)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    UNWIND [x IN row.yearlyRevenues WHERE NOT toString(x.year) IS NULL] AS yearlyRevenue\n",
    "    MERGE (s:`RevenueYear` {id: o.name + toString(yearlyRevenue.year)})\n",
    "    ON CREATE SET s.year = yearlyRevenue.year,\n",
    "                  s.revenue = toFloat(yearlyRevenue.revenue.amount),\n",
    "                  s.currency = yearlyRevenue.revenue.currency\n",
    "    MERGE (o)-[:HAS_YEARLY_REVENUE]->(s)\n",
    "    WITH s ORDER BY s.year\n",
    "    WITH collect(s) AS revenues\n",
    "    CALL apoc.nodes.link(revenues, \"NEXT_YEAR\")\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    UNWIND [x IN row.technographics WHERE NOT x.technology.name IS NULL] AS technographic\n",
    "    MERGE (s:`Technographic` {name: technographic.technology.name})\n",
    "    WITH * WHERE not s:Common\n",
    "    MERGE (o)-[:HAS_TECHNOGRAPHIC]->(s)\n",
    "    WITH s, technographic\n",
    "    UNWIND [x IN technographic.categories WHERE NOT x is null] AS category\n",
    "    MERGE (c:`TechnographicCategory` {name: category})\n",
    "    MERGE (s)-[:HAS_CATEGORY]->(c)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    UNWIND [x IN row.investments WHERE NOT x.date.str is null] AS investment\n",
    "    MERGE (is:`InvestmentSeries` {id: o.name + investment.date.str})\n",
    "    ON CREATE SET is.amount = investment.amount.value,\n",
    "                  is.series = investment.series,\n",
    "                  is.currency = investment.amount.currency,\n",
    "                  is.date = investment.dateClean\n",
    "    MERGE (o)-[:HAS_INVESTMENT]->(is)\n",
    "    WITH is, investment\n",
    "    UNWIND [x in investment.investors WHERE NOT x.diffbotUri is null] as investor\n",
    "    CALL apoc.merge.node(\n",
    "    [investor.type],\n",
    "    {id: investor.diffbotUri},\n",
    "    {name: investor.name, summary: investor.summary},\n",
    "    {}\n",
    "    ) YIELD node AS es\n",
    "    MERGE (es)-[:HAS_INVESTED]->(is)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    UNWIND [x in row.locations WHERE NOT x.address IS NULL] AS location\n",
    "    MERGE (s:`Place` {address: substring(location.address,0,8000)})\n",
    "    ON CREATE SET s.location = point({latitude:toFloat(location.latitude), longitude:toFloat(location.longitude)}),\n",
    "    s.city = location.city.name, s.region = location.region.name, s.country = location.country.name\n",
    "    MERGE (o)-[h:HAS_LOCATION]->(s)\n",
    "    SET h.isPrimary = location.isPrimary\n",
    "/*\n",
    "    WITH s, location\n",
    "    WHERE location.country IS NOT NULL AND location.country.diffbotUri IS NOT NULL\n",
    "    MERGE (c:`Country` {id: location.country.diffbotUri})\n",
    "    ON CREATE SET c.name = location.country.name\n",
    "    MERGE (s)-[hc:HAS_COUNTRY]->(c)\n",
    "*/\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "CALL (o, row) {\n",
    "//    WITH o, row\n",
    "    WITH o, row WHERE row.ceo IS NOT NULL and row.ceo.diffbotUri IS NOT NULL\n",
    "    MERGE (c:`Person` {id: row.ceo.diffbotUri})\n",
    "    ON CREATE SET c.name = row.ceo.name,\n",
    "                  c.summary = row.ceo.summary\n",
    "    MERGE (o)-[:HAS_CEO]->(c)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH o, row\n",
    "FOREACH (c IN [x in row.diffbotClassification WHERE x.version IS NOT NULL] |\n",
    "  MERGE (cl:Classification {id: c.version})\n",
    "  SET cl.name = c.name\n",
    "  MERGE (o)-[hc:HAS_CLASSIFICATION]->(cl)\n",
    "  SET hc.isPrimary = c.isPrimary\n",
    ")\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "F2z0DfbUU9gj",
    "outputId": "bd6c1b16-4e79-4a70-86f6-b0d3c6f6af90"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EagerResult(records=[], summary=<neo4j._work.summary.ResultSummary object at 0x177fd49d0>, keys=[])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "driver.execute_query(query, {\"data\": [el[\"entity\"] for el in data['data']]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "_kwlF5wDVmCT"
   },
   "outputs": [],
   "source": [
    "from typing import Dict, Any\n",
    "from urllib.parse import urlencode\n",
    "\n",
    "def process_entities(entity: str, type: str) -> Dict[str, Any]:\n",
    "    \"\"\"\n",
    "    Fetch relevant articles from Diffbot KG endpoint\n",
    "    \"\"\"\n",
    "    search_host = \"https://kg.diffbot.com/kg/v3/enhance?\"\n",
    "    params = {\"type\": type, \"id\": entity, \"token\": token}\n",
    "    encoded_query = urlencode(params)\n",
    "    url = f\"{search_host}{encoded_query}\"\n",
    "    return entity, requests.get(url).json()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, person_data = process_entities(\"E0Akwc0aKMKmTclQRNAiwRQ\", \"Person\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = person_data['data'][0]['entity']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_dates_person(data):\n",
    "    for entity in data['data']:\n",
    "        entity['entity']['birthDateClean'] = get_datetime(abs(entity['entity'].get('birthDate',{}).get('timestamp')))\n",
    "        for edu in entity['entity']['educations']:\n",
    "            edu['dateFrom'] = get_datetime(edu.get('from', {}).get('timestamp'))\n",
    "            edu['dateTo'] = get_datetime(edu.get('to',{}).get('timestamp'))\n",
    "        for emp in entity['entity']['employments']:\n",
    "            emp['dateFrom'] = get_datetime(emp.get('from', {}).get('timestamp'))\n",
    "            emp['dateTo'] = get_datetime(emp.get('to',{}).get('timestamp'))\n",
    "        for aw in entity['entity']['awards']:\n",
    "            aw['dateClean'] = get_datetime(aw.get('date', {}).get('timestamp'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "metadata": {},
   "outputs": [],
   "source": [
    "preprocess_dates_person(person_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'name': 'consultation', 'type': 'Miscellaneous'}]\n"
     ]
    }
   ],
   "source": [
    "print(t['interests'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    }
   ],
   "source": [
    "used = [\"diffbotUri\", \"description\", \"importance\", \"gender\", \"image\", \"wikipediaUri\", \"crunchbaseUri\", \"netWorth\", \"linkedInUri\",\n",
    "       \"birthPlace\", \"wikipediaPageviews\", \"birthDate\", \"nationalities\", \"languages\", \"educations\", \"birthDateClean\", \"employments\",\n",
    "       \"awards\", \"locations\", \"skills\", \"unions\", \"twitterUri\", \"youtubeUri\", \"height\", \"googlePlusUri\", \"instagramUri\", \"facebookUri\",\n",
    "       \"homepageUri\", \"children\", \"parents\", \"politicalAffiliation\", \"religion\", \"interests\"]\n",
    "ignore = [\"id\", 'origin', 'type', \"allNames\", \"wikipediaPageviewsLastYearGrowth\", \"origins\", \"crawlTimestamp\",\"nbOrigins\", \"summary\",\n",
    "         \"types\", \"images\", \"nbIncomingEdges\", \"wikipediaPageviewsLastQuarter\", \"wikipediaPageviewsLastQuarterGrowth\", \n",
    "          \"wikipediaPageviewsLastYear\", \"name\", \"allDescriptions\", \"wikipediaPageviewsGrowth\", \"allOriginHashes\", \"nameDetail\", \"age\",\n",
    "         \"allUris\", \"location\", \"nbFollowers\", \"weight\", \"nbFollowers\", \"emailAddresses\", \"phoneNumbers\"]\n",
    "available = [key for key in person_data['data'][0]['entity'].keys() if key not in used and key not in ignore]\n",
    "print(available)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EagerResult(records=[], summary=<neo4j._work.summary.ResultSummary object at 0x320b11a50>, keys=[])"
      ]
     },
     "execution_count": 174,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "driver.execute_query(\"CREATE CONSTRAINT IF NOT EXISTS FOR (n:Award) REQUIRE n.name IS UNIQUE;\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "metadata": {},
   "outputs": [],
   "source": [
    "person_query = \"\"\"\n",
    "UNWIND $data AS row\n",
    "MERGE (p:Person {id: row.diffbotUri})\n",
    "SET p.description = row.description,\n",
    "    p.importance = toFloat(row.importance),\n",
    "    p.gender = row.gender.normalizedValue,\n",
    "    p.wikipediaPageviews = toInteger(row.wikipediaPageviews),\n",
    "    p.summary = row.summary,\n",
    "    p.image = row.image,\n",
    "    p.crunchbaseUri = row.crunchbaseUri,\n",
    "    p.netWorth = toFloat(row.netWorth.value),\n",
    "    p.netWorthCurrency = row.netWorth.currency,\n",
    "    p.linkedInUri = row.linkedInUri,\n",
    "    p.birthCountry = row.birthPlace.country.name,\n",
    "    p.birthDate = row.birthDateClean,\n",
    "    p.nationalities = apoc.coll.toSet([el in row.nationalities | el.name]),\n",
    "    p.languages = apoc.coll.toSet([el in row.languages | el.str]),\n",
    "    p.skills = apoc.coll.toSet([el in row.skills | el.name]),\n",
    "    p.twitterUri = row.twitterUri,\n",
    "    p.youtubeUri = row.youtubeUri,\n",
    "    p.height = toFloat(row.height),\n",
    "    p.homepageUri = row.homepageUri,\n",
    "    p.instagramUri = row.instagramUri,\n",
    "    p.googlePlusUri = row.googlePlusUri,\n",
    "    p.facebookUri = row.facebookUri,\n",
    "    p.politicalAffiliation = apoc.coll.toSet([el in row.politicalAffiliation | el.name]),\n",
    "    p.religion = row.religion.str,\n",
    "    p.interests = apoc.coll.toSet([el in row.interests | el.name]),\n",
    "    p.name = row.name\n",
    "\n",
    "    \n",
    "WITH p, row\n",
    "CALL (p, row) {\n",
    "   UNWIND [x IN row.educations WHERE NOT x.institution.diffbotUri IS NULL] AS education\n",
    "   MERGE (o:Organization {id: education.institution.diffbotUri})\n",
    "   ON CREATE SET o.name = education.institution.name,\n",
    "                 o.summary = education.institution.summary\n",
    "   MERGE (p)-[h:HAS_EDUCATION_FROM]->(o)\n",
    "   SET h.dateFrom = education.dateFrom,\n",
    "       h.dateTo = education.dateTo,\n",
    "       h.isCurrent = education.isCurrent,\n",
    "       h.degree = education.degree.name\n",
    "  RETURN count(*) AS count\n",
    "}\n",
    "WITH p, row\n",
    "CALL (p, row) {\n",
    "   UNWIND [x IN row.employments WHERE NOT x.employer.diffbotUri IS NULL] AS employment\n",
    "   MERGE (o:Organization {id: employment.employer.diffbotUri})\n",
    "   ON CREATE SET o.name = employment.employer.name,\n",
    "                 o.summary = employment.employer.summary\n",
    "   MERGE (p)-[e:EMPLOYMENT]->(o)\n",
    "   SET e.isCurrent = employment.isCurrent,\n",
    "       e.dateFrom = employment.dateFrom,\n",
    "       e.dateTo = employment.dateTo,\n",
    "       e.title = employment.title,\n",
    "       e.categories = [el in employment.categories | el.name]\n",
    "   RETURN count(*) AS count\n",
    "}\n",
    "WITH p, row\n",
    "CALL (p, row) {\n",
    "   UNWIND [x IN row.awards WHERE NOT x.title IS NULL] AS award\n",
    "   MERGE (o:Award {name: award.title})\n",
    "   MERGE (p)-[e:RECEIVED_AWARD]->(o)\n",
    "   SET e.date = award.dateClean\n",
    "   RETURN count(*) AS count\n",
    "}\n",
    "WITH p, row\n",
    "CALL (p, row) {\n",
    "    UNWIND [x in row.locations WHERE NOT x.address IS NULL] AS location\n",
    "    MERGE (s:`Place` {address: substring(location.address,0,8000)})\n",
    "    ON CREATE SET s.location = point({latitude:toFloat(location.latitude), longitude:toFloat(location.longitude)}),\n",
    "    s.city = location.city.name, s.region = location.region.name, s.country = location.country.name\n",
    "    MERGE (p)-[h:HAS_LOCATION]->(s)\n",
    "    SET h.isPrimary = location.isPrimary\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH p, row\n",
    "CALL (p, row) {\n",
    "    UNWIND [x in row.unions WHERE NOT x.person.diffbotUri IS NULL] AS union\n",
    "    MERGE (s:`Person` {id: union.person.diffbotUri})\n",
    "    ON CREATE SET s.name = union.person.name,\n",
    "                  s.summary = union.person.summary\n",
    "    MERGE (p)-[u:UNION]-(s)\n",
    "    SET u.type = union.type\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH p, row\n",
    "CALL (p, row) {\n",
    "    UNWIND [x in row.children WHERE NOT x.diffbotUri IS NULL] AS child\n",
    "    MERGE (s:`Person` {id: child.diffbotUri})\n",
    "    ON CREATE SET s.name = child.name,\n",
    "                  s.summary = child.summary\n",
    "    MERGE (p)<-[:HAS_PARENT]-(s)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "WITH p, row\n",
    "CALL (p, row) {\n",
    "    UNWIND [x in row.parents WHERE NOT x.diffbotUri IS NULL] AS parent\n",
    "    MERGE (s:`Person` {id: parent.diffbotUri})\n",
    "    ON CREATE SET s.name = parent.name,\n",
    "                  s.summary = parent.summary\n",
    "    MERGE (p)-[:HAS_PARENT]->(s)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "CALL (p, row) {\n",
    "    UNWIND [x in row.siblings WHERE NOT x.diffbotUri IS NULL] AS sibling\n",
    "    MERGE (s:`Person` {id: sibling.diffbotUri})\n",
    "    ON CREATE SET s.name = sibling.name,\n",
    "                  s.summary = sibling.summary\n",
    "    MERGE (p)-[:SIBLING]->(s)\n",
    "    RETURN count(*) AS count\n",
    "}\n",
    "RETURN count(*)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EagerResult(records=[<Record count(*)=1>], summary=<neo4j._work.summary.ResultSummary object at 0x320d18050>, keys=['count(*)'])"
      ]
     },
     "execution_count": 184,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "driver.execute_query(person_query, {\"data\": [el[\"entity\"] for el in person_data['data']]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyNr0ll91B76jYwbS8sA0Egd",
   "include_colab_link": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
